import math
import numpy as np
from core.cuda import cuda, Cuda


class DataLoader:
    """
    用于加载数据的数据加载器类。
    """

    def __init__(self, dataset, batch_size: int, shuffle: bool = True, gpu: bool = False):
        """
        初始化数据加载器。

        Args:
            dataset: 数据集，可迭代对象。
            batch_size (int): 批处理大小。
            shuffle (bool, optional): 是否在每个迭代周期中随机打乱数据。默认为True。
            gpu (bool, optional): 是否使用GPU。默认为False。
        """
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.data_size = len(dataset)
        self.max_iter = math.ceil(self.data_size / batch_size)
        self.gpu = gpu

        self.reset()

    def reset(self):
        """
        为下一个 epoch 重置迭代器的状态。

        如果 `shuffle` 为 True，则随机排列数据集索引，否则使用顺序索引。
        """
        self.iteration = 0  # 重置迭代计数器
        # 如果 shuffle 为 True，则使用随机排列，否则使用顺序索引
        self.index = np.random.permutation(len(self.dataset)) if self.shuffle else np.arange(len(self.dataset))

    def __iter__(self):
        return self

    # 迭代器
    def __next__(self):
        """
        从数据集中获取下一批数据。

        Raises:
            StopIteration: 当所有批次都被处理时。

        Returns:
            tuple: 包含输入数据和相应标签的元组。
        """
        if self.iteration >= self.max_iter:
            self.reset()
            raise StopIteration

        i, batch_size = self.iteration, self.batch_size
        batch_index = self.index[i * batch_size:(i + 1) * batch_size]
        batch = [self.dataset[i] for i in batch_index]

        xp = Cuda(self.gpu).module
        # 从批次中提取输入数据和标签
        x = xp.array([example[0] for example in batch])
        t = xp.array([example[1] for example in batch])

        self.iteration += 1
        return x, t

    def next(self):
        return self.__next__()

    def to_cpu(self):
        self.gpu = False

    def to_gpu(self):
        self.gpu = True
